#ifndef HB_CNN_H
#define HB_CNN_H

#include "tensor.h"
#include "gradient.h"
#include "fc.h"
#include "pool.h"
#include "relu.h"
#include "conv.h"
#include "drop_out.h"
using namespace std;

class Model
{
public:
	Model();

	void add_conv( uint16_t stride, uint16_t extend_filter, uint16_t number_filters, tdsize in_size );

	void add_relu( tdsize in_size );

	void add_pool( uint16_t stride, uint16_t extend_filter, tdsize in_size );

	void add_fc( tdsize in_size, int out_size );

	tdsize& output_size();

	int predict();

	tensor_t<float>& predict_info();

	void forward( tensor_t<float>& data );

	float train( tensor_t<float>& data, tensor_t<float>& label );

private:
	vector<layer_t*> layers;
};

#endif //HB_CNN_H